import numpy as np
import os
from PIL import Image
import glob
from yolov3_by_myself.get_xml import PascalVocXmlParser


def get_parse(ann_fname, input_size):
    parser = PascalVocXmlParser()
    fname = parser.get_fname(ann_fname)
    weight = parser.get_width(ann_fname)
    height = parser.get_height(ann_fname)
    labels = parser.get_labels(ann_fname)
    boxes = parser.get_boxes(ann_fname)

    for i in range(len(boxes)):
        # 对boxes进行归一化
        boxes[i][0] = boxes[i][0] / weight * input_size
        boxes[i][1] = boxes[i][1] / weight * input_size
        boxes[i][2] = boxes[i][2] / height * input_size
        boxes[i][3] = boxes[i][3] / height * input_size

    return fname, labels, boxes


def get_IOU(box1, box2):
    w_min = min(box1[1], box2[1])
    h_min = min(box1[3], box2[3])
    w = w_min - box1[0]
    h = h_min - box1[2]

    intersect = w * h
    merge = (box1[1] - box1[0]) * (box1[3] - box1[2]) + (box2[1] - box2[0]) * (box2[3] - box2[2])
    IOU = intersect / (merge - intersect)

    return IOU


def get_anchor(anchors, box):
    IOUlist = []
    anchorslist = np.zeros((len(anchors), 4), dtypa='float32')
    for i in range(len(anchorslist)):
        anchorslist[i][0] = box[0]
        anchorslist[i][1] = anchors[i][0] + anchorslist[i][0]
        anchorslist[i][2] = box[2]
        anchorslist[i][3] = anchors[i][1] + anchorslist[i][2]

        IOU = get_IOU(box, anchorslist[i])
        IOUlist.append(IOU)

    anchor = IOUlist.index((max(IOUlist)))
    return anchor


def get_ytrue(boxes, anchors, anchor_shape, b, pattern_shape, input_size, classes, labels, ytrues):
    newbox = np.zeros((4), dtype='float32')
    for i in range(len(boxes)):
        anchor = get_anchor(anchors, boxes[i])

        layer_anchor = anchor // anchor_shape[1]
        box_anchor = anchor % anchor_shape[1]

        rate = pattern_shape[layer_anchor] / input_size

        cent_x = (boxes[i][0] + boxes[i][1]) / 2 * rate
        cent_y = (boxes[i][2] + boxes[i][3]) / 2 * rate

        x = np.floor(cent_x).astype('int32')
        y = np.floor(cent_y).astype('int32')
        w = boxes[i][1] - boxes[i][0]
        h = boxes[i][3] - boxes[i][2]

        c = classes.index(labels[i])

        newbox[0] = cent_x - x
        newbox[1] = cent_y - y
        newbox[2] = np.log(w / anchors[anchor][0])
        newbox[3] = np.log(h / anchors[anchor][1])

        ytrues[layer_anchor][b, x, y, box_anchor, 0:4] = newbox[0:4]
        ytrues[layer_anchor][b, x, y, box_anchor, 4] = 1
        ytrues[layer_anchor][b, x, y, box_anchor, 5 + c] = 1

    return ytrues


def get_img(img_dir,fname,input_size):
    img_fname = os.path.join(img_dir,fname)
    image = Image.open(img_fname)
    img = np.array(image.resize((input_size,input_size)))
    return img


def generator(batch_size, classes, ann_fnames, input_size, anchors,img_dir):
    pattern_shape = [52, 32, 16]
    anchor_shape = [3, 3]
    n = len(ann_fnames)
    i = 0
    while True:
        inputs = []
        ytrues = [np.zeros(batch_size, pattern_shape[l], pattern_shape[l], anchor_shape[1],
                           5 + len(classes)) for l in range(anchor_shape[0])]

        for b in range(batch_size):
            if i == 0:
                np.random.shuffle(ann_fnames)

            fname, labels, boxes = get_parse(ann_fnames[i], input_size)
            ytrues = get_ytrue(boxes, anchors, anchor_shape, b, pattern_shape, input_size, classes,
                               labels, ytrues)
            img = get_img(img_dir,fname,input_size)
            inputs.append(img)
            i = (i + 1) % n
        inputs = np.array(inputs)

        yield inputs, ytrues


def main():
    root = os.path.dirname(__file__)
    ann_dir = os.path.join(root, "data", "ann", "*.xml")
    ann_fnames = glob.glob(ann_dir)
    img_dir = os.path.join(root, "data", "img")
    anchors = np.array([[8,73],[8,24],[13,32],[19,51],[35,64],[25,37],[22,164],[95,195],[95,104]])
    batch_size = 1
    input_size = 416
    classes = ['Car','Person','Cyclist','Truck','Van']
    for  inputs,ytrues in generator(batch_size,classes,ann_fnames,input_size,anchors,img_dir):
        print(input.shape,ytrues[0].shape,ytrues[1].shape,ytrues[2].shape)


if __name__ == '__main__':
    main()

